# -*- coding: utf-8 -*-
# @Time    : 2024/5/10 17:16
# @Author  : Z.J. Zhang
# @Email   : zijingzhang@mail.ustc.edu.cn
# @File    : doppler_force.py
# @Software: PyCharm
import typing

import matplotlib.pyplot as plt
import numpy
import scipy.constants as C

import mot.field.magnetic
from mot.atom import TransitionLineData, Rb87_D2
from mot.field.base import RectGrid3D
from mot.field.laser import LaserFieldGenerator, NonDiffractionLaserBeam, LaserInfo, FieldNames
from mot.field.laser import set_frequency
from mot.field.magnetic import BGenerators
from mot.field.magnetic import CurrentCarryingLoop
from mot.gmot.grating import Sector, GratingSection, LaserDiffractedByGratingSection
from myutils.geometrical import Plane


class MOT:
    def __init__(self, grid: RectGrid3D,
                 B_generator: BGenerators,
                 laser_field_generators: typing.List[LaserFieldGenerator],
                 transition_line_data: TransitionLineData = Rb87_D2):
        self.grid = grid
        self.B_generator = B_generator
        self.laser_field_generators = laser_field_generators
        self.laser_field_names = FieldNames([ls.laser_info.name for ls in self.laser_field_generators])
        self._dict_laser_name_to_index = {ls.laser_info.name: i for i, ls in enumerate(self.laser_field_generators)}

        self.transition_line_data = transition_line_data
        self.cache()

    def cache(self, ):
        grid = self.grid
        self.B_generator.cache(grid)
        B_data: numpy.ndarray = grid.V[..., grid.get_index_of_parameter_name([
            mot.field.magnetic.FieldNames.Bx.name,
            mot.field.magnetic.FieldNames.By.name,
            mot.field.magnetic.FieldNames.Bz.name,
        ])]

        for lfg in self.laser_field_generators:
            lfg.cache(grid)
            # 波矢与磁场的夹角的余弦值, 计算方式：
            # cos(<a,b>) = dot(a,b) / (||a|| * ||b||)
            costheta = (numpy.matrix(B_data.reshape((-1, 3))) * numpy.matrix(
                lfg.laser_info.k.reshape((3, 1)))).A.reshape(
                B_data.shape[:-1]) / (
                               numpy.linalg.norm(B_data, axis=B_data.ndim - 1) * numpy.linalg.norm(lfg.laser_info.k))

            grid.set_parameter_data({
                lfg.laser_info.field_name_costheta: costheta,
            })
        grid.set_parameter_data({
            self.laser_field_names.sumI: grid.V[..., grid.get_index_of_parameter_name(
                self.laser_field_names.Is)].sum(axis=grid.V.ndim - 1)
        })

        # return grid

    @staticmethod
    def eta(s0, s, costheta):
        """
        相对跃迁强度

        relative transition strength

        Ref: Maximized atom number for a grating magneto-optical trap via machine-learning assisted parameter optimization

        See explanation of Eq. (1)

        :param s0: polarization of the beam, s0 = +1 means right polarization
        :param s:
        :return:
        """
        if s != 0:
            return (1 - s * s0 * costheta) ** 2 / 4
            # return (1 + s * s0 * costheta) ** 2 / 4
        else:
            return (1 - costheta ** 2) / 2
        # return  (lambda s0,costheta :numpy.piecewise(s,
        #                       [s == 0, ], [
        #                           (1 - costheta ** 2) / 2,
        #                           (1 - s * s0 * costheta) ** 2 / 4
        #                       ]))(s0, costheta)

    def scattering_force_from_laser(
            self,
            x: numpy.ndarray, y: numpy.ndarray, z: numpy.ndarray,
            vx: numpy.ndarray, vy: numpy.ndarray, vz: numpy.ndarray,
    ):
        costhetas_Is_sumI = self.grid.get(x, y, z, [*(self.laser_field_names.costhetas + self.laser_field_names.Is),
                                                    self.laser_field_names.sumI, self.B_generator.field_names.B.name])
        costhetas = costhetas_Is_sumI[..., :len(self.laser_field_names.costhetas)]
        Is = costhetas_Is_sumI[..., len(self.laser_field_names.costhetas):len(self.laser_field_names.costhetas) + len(
            self.laser_field_names.Is)]
        sumI = costhetas_Is_sumI[..., -2]
        B = costhetas_Is_sumI[..., -1]
        Isat = self.transition_line_data.I_sat

        v = numpy.array([vx, vy, vz])
        _sum = 0
        for i, laser_i in enumerate(self.laser_field_generators):
            for s in [-1, 0, 1]:
                delta = laser_i.laser_info.detuning_Hz - (
                        numpy.matrix(laser_i.laser_info.k) * numpy.matrix(v.reshape((3, -1)))).A.reshape(
                    vx.shape) - s * self.transition_line_data.mu_F * B / C.hbar
                _sum += laser_i.laser_info.k * (self.eta(
                    laser_i.laser_info.polarization.value,
                    s, costhetas[..., i]
                ) * Is[..., i] / Isat / (1 + sumI / Isat + (2 * delta / self.transition_line_data.Gamma) ** 2)).reshape(
                    (*vx.shape, 1))

        return C.hbar * self.transition_line_data.Gamma / 2 * _sum


if __name__ == '__main__':
    pass